package com.slipper.ai.vector;

import com.slipper.ai.model.embedding.EmbeddingModelFactory;
import com.slipper.constants.Constants;
import com.slipper.core.gson.GsonUtil;
import com.slipper.system.ai.bo.AiModelBO;
import com.slipper.system.ai.bo.KnowledgeBO;
import com.slipper.system.ai.bo.KnowledgeEmbeddingBO;
import com.slipper.system.ai.service.AiModelService;
import com.slipper.system.ai.service.KnowledgeEmbeddingService;
import com.slipper.system.ai.service.KnowledgeService;
import jakarta.annotation.Resource;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.stream.Collectors;

@Component
public class KnowledgeVectorService {
    @Resource
    private AiModelService aiModelService;
    @Resource
    private KnowledgeService knowledgeService;
    @Resource
    private KnowledgeEmbeddingService knowledgeEmbeddingService;

    public Vector loadByKnowledgeId(String knowledgeId, AiModelBO model) {
        List<KnowledgeEmbeddingBO> dataList = knowledgeEmbeddingService.queryListByKnowledgeId(knowledgeId);
        Vector vector = new Vector(EmbeddingModelFactory.create(model));
        dataList.forEach(data -> vector.add(
                data.getId(),
                data.getContent(),
                GsonUtil.toMap(data.getMetadata()),
                GsonUtil.fromJson(data.getEmbedding(), float[].class)));
        return vector;
    }

    public String similarityText(String knowledgeId, String query) {
        return similarityText(knowledgeId, query, null, null);
    }

    public String similarityText(String knowledgeId, String query, Integer topSize, Double threshold) {
        return similaritySearch(knowledgeId, query, topSize, threshold).stream()
                .map(Document::getText)
                .collect(Collectors.joining("\n\n"));
    }

    public List<Document> similaritySearch(String knowledgeId, String query) {
        return similaritySearch(knowledgeId, query, null, null);
    }

    public List<Document> similaritySearch(String knowledgeId, String query, Integer topSize, Double threshold) {
        KnowledgeBO knowledge = knowledgeService.queryBy(knowledgeId);
        AiModelBO model = aiModelService.queryBy(knowledge.getModelId());

        SearchRequest.Builder builder = SearchRequest.builder();
        if (topSize != null) {
            builder.topK(topSize);
        } else if (knowledge.getTopSize() != null) {
            builder.topK(knowledge.getTopSize());
        }
        if (threshold != null) {
            builder.similarityThreshold(threshold);
        } else if (Constants.YES.equals(knowledge.getUseSimilarityThreshold())
                && knowledge.getSimilarityThreshold() != null) {
            builder.similarityThreshold(knowledge.getSimilarityThreshold().doubleValue());
        }
        builder.query(query);
        SearchRequest request = builder.build();

        Vector vector = loadByKnowledgeId(knowledgeId, model);
        return vector.similaritySearch(request);
    }
}
